import torch
from torch.utils.data import Dataset, DataLoader
import pandas as pd
import config

class TranslationDataset(Dataset):
    def __init__(self, data_path):
        self.data_path = data_path
        self.data = pd.read_json(self.data_path, orient='records', lines=True).to_dict(orient='records')

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        input_tensor = torch.tensor(self.data[index]['zh'], dtype=torch.long)
        target_tensor = torch.tensor(self.data[index]['en'], dtype=torch.long)
        return input_tensor, target_tensor


def get_dataloader(train = True):
    data_path = config.PROCESSED_DIR / 'indexed_train.jsonl' if train else config.PROCESSED_DIR / 'indexed_test.jsonl'
    dataset = TranslationDataset(data_path)
    return DataLoader(dataset,batch_size=config.BATCH_SIZE,shuffle=train)

if __name__ == '__main__':

    loader = get_dataloader(train=True)

    iterator = iter(loader)
    input,target = next(iterator)
    print(input.shape)
    print(target.shape)